#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""等高线转DEM
通过调用 gdal_rasterize gdal_fillnodata.py 这两个命令行工具, 完成等高线到 DEM 的转换
思路是首先将等高线矢量转栅格然后填充 NoData 值。

使用过程中必须指定目标栅格的 *像元大小* , 包含 x 方向和 y 方向, 注意这两个值带引号, 用空格分开。
此外,需要特别注意的参数有
    [-a 属性名] shapefile 文件中作为高度的字段名, 默认 Elevation
    [-md 最大距离] 在填充 NoData 过程中距离有值的像素点周围的最远像素点距离,该值设置过大可能会影响结果准确度,太小会在生成的 tif 文件中留下空白区域, 默认 16

Example:
    python contour2dem.py -tr "0.00034 0.00034" -md 24 contour_line.shp out.tif
"""

import os, sys, uuid, tempfile
import shlex, subprocess
from osgeo import ogr

def main(shape_fn, raster_fn, tr, an, md):
    datasource = ogr.Open(shape_fn)
    layer = datasource.GetLayer()
    min_x, max_x, min_y, max_y = layer.GetExtent()

    rasterize_fn = f"{uuid.uuid4()}.tif"
    rasterize_fn = os.path.join(tempfile.gettempdir(), rasterize_fn)

    cmd_rasterize = f"gdal_rasterize -l {layer.GetName()} -a {an} -tr {tr} " +\
                    f"-te {min_x} {min_y} {max_x} {max_y} -a_nodata 0.0 {shape_fn} {rasterize_fn}"
    args_rasterize = shlex.split(cmd_rasterize)

    rasterize_result = subprocess.run(args_rasterize, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stderr = rasterize_result.stderr.decode('utf-8')
    if stderr != '':
        print(f"error: {stderr}")
        sys.exit(1)

    cmd_fill = f"gdal_fillnodata.py -md {md} -b 1 -of GTiff " +\
               f"{rasterize_fn} {raster_fn}"
    args_fill = shlex.split(cmd_fill)

    fill_result = subprocess.run(args_fill, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stderr = fill_result.stderr.decode('utf-8')
    if stderr != '':
        print(f"error: {stderr}")
        sys.exit(1)

    print("process complete!")

if __name__ == "__main__":
    shp_fn = None
    raster_fn = None
    target_resolution = None
    attribute_name = None
    max_distance = None

    argv = sys.argv

    i = 1
    while i < len(argv):
        arg = argv[i]

        if arg == '-tr':
            i = i + 1
            target_resolution = argv[i]

        elif arg == '-a':
            i = i + 1
            attribute_name = argv[i]

        elif arg == '-md':
            i = i + 1
            max_distance = argv[i]

        elif shp_fn is None:
            shp_fn = argv[i]

        elif raster_fn is None:
            raster_fn = argv[i]

        else:
            print("python script.py [-tr xres yres] [-md max_distance] shape.shp out.tif")
            sys.exit(1)

        i = i + 1

    if target_resolution is None:
        print("need target resolution")
        sys.exit(1)

    if max_distance is None:
        max_distance = 16

    if attribute_name is None:
        attribute_name = 'Elevation'

    if shp_fn is None:
        print("need shapefile")
        sys.exit(1)

    if raster_fn is None:
        raster_fn = 'OUT.tif'

    main(shp_fn, raster_fn, target_resolution, attribute_name, max_distance)

